"""
实现语义搜索
本实践案例通过文本向量模型获取产品描述对应的向量，然后基于向量，计算评论与产品描述之间的余弦相似度，最终返回与产品描述最相关的评论。
"""
"""
1、导入库包
导入所需的库包，并设置API Key，为后续的数据处理和分析做准备。
"""
import os
import pandas as pd
import numpy as np
from ast import literal_eval
import dashscope
from dashscope import TextEmbedding

# 设置DashScope的API Key
dashscope.api_key = os.getenv("DASHSCOPE_API_KEY")
"""
2、读取CSV文件并处理嵌入数据
读取CSV文件fine_food_reviews.csv（其中包含评论和评论对应的向量），转换embedding列的字符串数据为NumPy数组以便后续处理。
"""
# 读取CSV数据
datafile_path = "fine_food_reviews.csv"
df = pd.read_csv(datafile_path)

# 将embedding列的字符串数据转换为NumPy数组
df["embedding"] = df.embedding.apply(literal_eval).apply(np.array)
"""
3、定义文本嵌入的函数
该函数将文本转换为向量，并根据输入类型的不同返回单个向量或向量列表。
"""
# 定义文本嵌入的函数
def generate_embeddings(text):
    rsp = TextEmbedding.call(model=TextEmbedding.Models.text_embedding_v1, input=text)
    embeddings = [record['embedding'] for record in rsp.output['embeddings']]
    return embeddings if isinstance(text, list) else embeddings[0]
"""
4、定义计算余弦相似度函数和搜索评论函数
利用余弦相似度来度量文本之间的相似性，实现了一个基于内容的评论搜索功能。用户可以通过输入产品描述，检索与该产品描述最相关的评论。
"""
# 定义计算余弦相似度函数
def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

# 搜索特定产品的评论
def search_reviews(df, product_description, n=3, pprint=True):
    product_embedding = generate_embeddings(product_description)
    df["similarity"] = df.embedding.apply(lambda x: cosine_similarity(x, product_embedding))
    results = (
        df.sort_values("similarity", ascending=False)
        .head(n)
        .combined.str.replace("Title: ", "")
        .str.replace("; Content:", ": ")
    )
    if pprint:
        for r in results:
            print(r[:200])
            print()
    return results
"""
5、测试搜索功能
测试搜索功能，调用search_reviews函数查找与宠物食品相关的评论，并返回2条相似评论。示例代码如下：
"""
results = search_reviews(df, "pet food", n=2)
# results = search_reviews(df, "delicious beans", n=3)
# results = search_reviews(df, "whole wheat pasta", n=3)
# results = search_reviews(df, "bad delivery", n=1)
# results = search_reviews(df, "spoilt", n=1)